#include "ExecutionEngine/fpga_sim_api.h"

#include <dirent.h>
#include <sys/stat.h>

//===----------------------------------------------------------------------===//
// Create simulation files for simulation environment to call
//===----------------------------------------------------------------------===//
bool CreateMd5SimFile(const string &sim_template_path,
                      const string &output_file_path,
                      const vector<unsigned char> &data,
                      vector<uint32_t> &result) {
  ifstream template_file(sim_template_path);
  if (!template_file) {
    cerr << "Error opening file: " << sim_template_path << endl;
    return false;
  }
  ofstream output_file(output_file_path);
  if (!output_file) {
    cerr << "Error opening file: " << output_file_path << endl;
    return false;
  }

  string line;
  bool found_end_module = false;
  while (getline(template_file, line)) {
    if (!found_end_module && line.find("endmodule") != string::npos) {
      output_file
          << "always @(start_setinputdata) begin\n    @(posedge clk) begin\n";
      for (size_t i = 0; i < data.size(); i += 4) {
        uint32_t val = (static_cast<uint32_t>(data[i]) << 24) |
                       (static_cast<uint32_t>(data[i + 1]) << 16) |
                       (static_cast<uint32_t>(data[i + 2]) << 8) |
                       static_cast<uint32_t>(data[i + 3]);
        output_file << "        input" << dec << i / 4 << " <= 32'h" << hex
                    << setw(8) << setfill('0') << val << ";\n";
      }
      // add initial value to input16-input19
      for (size_t i = 64, j = 0; i < 80 && j < 4; i += 4, ++j) {
        output_file << "        input" << dec << i / 4 << " <= 32'h" << hex
                    << setw(8) << setfill('0') << result[j] << ";\n";
      }
      output_file << "    end\nend\n\n";
      found_end_module = true;
    }
    output_file << line << "\n";
  }
  return true;
}

//===----------------------------------------------------------------------===//
// Perform simulation and obtain simulation results
//===----------------------------------------------------------------------===//
bool GetMd5ResultFromVCS(const string &src_files_path,
                         const string &sim_file_path,
                         vector<uint32_t> &result) {
  regex pattern(
      R"(^a = ([a-f0-9]{8}), b = ([a-f0-9]{8}), c = ([a-f0-9]{8});  d = ([a-f0-9]{8}))");
  string command = "vcs -R -full64 -debug_access+none -nolog " +
                   src_files_path + " " + sim_file_path;
  cout << "compute md5..." << endl;
  FILE *vcs_proc = popen(command.c_str(), "r");
  if (!vcs_proc) {
    cerr << "Error running VCS process" << endl;
    return false;
  }

  char buffer[1024];
  smatch match;
  string line;
  bool found_output = false;
  while (fgets(buffer, sizeof(buffer), vcs_proc) != nullptr) {
    line = buffer;
    if (regex_search(line, match, pattern)) {
      found_output = true;
      for (size_t i = 1; i < 5; ++i) {
        // result.push_back(stoul(match[i].str(), nullptr, 16));
        uint32_t tmp = stoul(match[i].str(), nullptr, 16);
        result[i - 1] = result[i - 1] + tmp;
      }
      break;
    }
  }
  pclose(vcs_proc);
  return found_output;
}

//===----------------------------------------------------------------------===//
// Plaintext filling algorithm
//===----------------------------------------------------------------------===//
vector<unsigned char> md5_padding(const vector<unsigned char> &input_data) {
  size_t padding_len = 64 - ((input_data.size() + 9) % 64);
  size_t output_length = padding_len + input_data.size() + 9;
  vector<unsigned char> output_data(output_length, 0);
  memcpy(output_data.data(), input_data.data(), input_data.size());
  output_data[input_data.size()] = 0x80;
  uint64_t bit_length = input_data.size() * 8;
  // md5 little-endian
  for (int i = 0; i < 8; ++i) {
    output_data[output_length - 8 + i] = (bit_length >> (8 * i)) & 0xFF;
  }
  // Input14 swap value
  swap(output_data[56], output_data[59]);
  swap(output_data[57], output_data[58]);
  // Input15 swap value
  swap(output_data[60], output_data[63]);
  swap(output_data[61], output_data[62]);
  return output_data;
}

bool md5_padded(const string &sim_template_path, const string &output_file_path,
                const string &src_files_path,
                const vector<unsigned char> &input_data,
                vector<uint32_t> &result) {
  if (!CreateMd5SimFile(sim_template_path, output_file_path, input_data,
                        result))
    return false;
  if (!GetMd5ResultFromVCS(src_files_path, output_file_path, result)) {
    cerr << "Error getting result from VCS!" << endl;
    return false;
  }
  return true;
}

void computeMd5(const string &sim_template_path, const string &output_file_path,
                const string &src_files_path,
                const vector<unsigned char> &input_data, uint8_t *cipher_addr) {
  vector<vector<unsigned char>> split_input;
  vector<uint32_t> compute_result = {0x67452301, 0xEFCDAB89, 0x98BADCFE,
                                     0x10325476};

  // 1. padding plaintext
  vector<unsigned char> padded_data = md5_padding(input_data);
  // 2. split plaintexted by pad, 512 bits per section
  for (size_t i = 0; i < padded_data.size(); i += 64) {
    size_t end = min(i + 64, padded_data.size());
    split_input.emplace_back(padded_data.begin() + i,
                             padded_data.begin() + end);
  }
  // comput md5
  for (size_t i = 0; i < split_input.size(); ++i) {
    if (!md5_padded(sim_template_path, output_file_path, src_files_path,
                    split_input[i], compute_result)) {
      break;
    }
  }
  if (!cipher_addr) {
    cerr << "Error cipher_addr is NULL!" << endl;
  } else {
    memcpy(cipher_addr, compute_result.data(), 20);
  }
}

//===----------------------------------------------------------------------===//
// MD5 encryption interface
//===----------------------------------------------------------------------===//
void fpga_md5(uint8_t *plain_addr, uint32_t plain_len, uint8_t *cipher_addr) {

  // Get environment variables
  string env_name = string(MD5_SIM_PATH);
  const char *algo_path = getenv(env_name.c_str());
  if (algo_path == nullptr) {
    cerr << "Not set env variable " << env_name << "." << endl;
    exit(1);
  }
  // template.v
  string sim_template_path = string(algo_path) + string(MD5_SIM_TMP_FILE);
  // *.v  verilog file
  string src_files_path = string(algo_path) + string(MD5_SIM_SRC_FILE);
  // testbentch.v
  string output_file_path = string(algo_path) + string(MD5_SIM_O_FILE);

  vector<unsigned char> input_data(plain_addr, plain_addr + plain_len);
  // compute MD5
  computeMd5(sim_template_path, output_file_path, src_files_path, input_data,
             cipher_addr);
}
